function [optima, RegionSet, SampleSet, RegionSampleId, OptSetRemain] = prsmmo_test( Problem, AlgorithmM, OptSet )
% prsmmo: Partition-based Random Search for Multimodal Optimization for 
%   minimization problems.
% decision variable: d dimensions.
% feature of regions: m dimensions.
%
%INPUT
%   Problem.
%       Dimension: scalar, the dimension of the decision variables
%       Domain: 1 -by- m vector, features of the domain
%       Partition: function handle of partitioning strategy
%           [RegionNum, Region, RegionSampleId] = Partition(Region, SampleSet)
%               RegionNum: scalar, number of new partitioned regions
%               Region: RegionNum -by- (2+m) matrix,
%                   [partition depth, partitionable, features ...]
%                   partition depth = 0 for the original domain
%                   partitionable: scalar, 1: partitionable, 0: non-partitionable
%               RegionSampleId: RegionNum -by- 1 cell, the sample id that
%                   belongs to each region
%               SampleSet: [SampleID, objective value, x1, x2, ...]
%       Sampling: function handle of sampling in a particular region
%           [sample] = Sampling(n, region)
%               n: number of samples
%               region: 1 -by- m vector, [features...]
%               sample: n -by- (1+d) matrix, [objective value, x1, x2, ...]
%   AlgorithmM.
%       n0: scalar, the base sample size
%       QuantileLevel: scalar, <0.5, the quantile level of promisinhg region definition
%       NewBudget: scalar, added budget size at each stage
%       MaxSampleSize: scalar, maximal sample size in one region
%       Radius: scalar, the radius distance that be regarded as neighbor
%       StopCriteria: [X,Y] (the maximal budget size is 1e6)
%           [1,Y]: reach maximum budget Y
%           [2,Y]: reach the required number of optima Y
%           [3,Y]: the number of optima is not changed in Y iterations
%
%OUTPUT
%   optima: ~ -by- (2+d) matrix, the final optimal solution set,
%         [sample id, objective value, x1, x2, ...]
%   RegionSet: ~ -by- (2+m) matrix, the whole subregion set
%         [partition depth, partitionable, features ...]
%   SampleSet: ~ -by- (1+d) matrix, the original sample set
%         [objective value, x1, x2, ...]
%   RegionSampleId: ~ -by- 1 cell, the id of the samples in each region

%% INITIALIZATION
% problem setting
d = Problem.Dimension;
domain = [0, 1, Problem.Domain]; % [partition depth, partitionable, features ...]
Partition = Problem.Partition;
Sampling = Problem.Sampling;
% algorithm setting
n0 = AlgorithmM.n0;
alpha = AlgorithmM.QuantileLevel;
Delta = AlgorithmM.NewBudget;
MaxSampleSize = AlgorithmM.MaxSampleSize;
StopCriteria = AlgorithmM.StopCriteria;
N = (StopCriteria(1)==1)*StopCriteria(2) + (StopCriteria(1)~=1)*1e6; % the maximal budget size
radius = AlgorithmM.radius;
% the optimal set for testing purpose
if ~isempty(OptSet)
    OptSetRemain = OptSet.sol;
end
% algorithm archive
RegionSet = zeros(round(N/n0), length(domain));  % [partition depth, partitionable, features ...]
RegionSet(1,:) = domain;
RegionsNum = 1;  % the number of exist regions
PartitionableNum = 1;  % the number of exist partitionable regions
RegionSampleId = cell(round(N/n0), 1);  % id of group samples
GroupInf = zeros(round(N/n0), 5); % [size, mean, std, adjusted size, weight]
SampleSet = zeros(N, 2+d);  % [SampleID, objective value, x1, x2, ...]
SampleSet(:,1) = 1:size(SampleSet,1);
N0 = 0;  % budget allocated
OptCandidate = [zeros(N,1),ones(N,1)]; % [Solution Id, 1: not be dominated by neibours]
OptCandidateNum = 0;
OptCandidate_checked = 0;  % number of optimal candidates that have been checked by the optima set
%% OPTIMIZATION
NumOpt = 0;  % the current optimal value
NumIterationRepeat = 0;  % the number of iterations that the number of optima is not changed
PartitionList = zeros(Delta,1);  % the list of regions to be partitioined
PartitionList(1) = 1;
PartitionListNum = 1;  % the number of the partition list
CurrentMaxPartitionDepth = 0;  % the maximum depth of partitioning at current iteration
WeightChangedList = zeros(round(N/n0),1);  % the list of regions whose weight is changed
WeightChangedListNum = 0;  % the number of the region list whose weight is changed
all_weight_changed = 0;  % if all weights should be recalculated
sum_adj_size = 0;  % the sum of the adjusted size
sum_weight = 0;  % the sum of the weights
NewOptCandidate = 0;  % if there are new regions reach the non-partitionable size
BestGroupInf = [1,inf,0];  % the best group information
%       [group id, quantile, modified sample size]
optima = [];  % the set of optima
% subsequent iterations
while N0 < N && PartitionableNum>0 ...
    && (StopCriteria(1)~=2 || NumOpt<StopCriteria(2)) ...
    && (StopCriteria(1)~=3 || NumOpt==0 || NumIterationRepeat<StopCriteria(2))...
    && (isempty(OptSet) || ~isempty(OptSetRemain))
%% partition partitionable regions
    i = 1;
    while i <= PartitionListNum
        % partition
        iPartitionRegion = PartitionList(i);
        [NewRegionNum, NewRegion, NewRegionSampleId]...
            = Partition(RegionSet(iPartitionRegion,:), SampleSet(RegionSampleId{iPartitionRegion},:));
        NewRegionId = [iPartitionRegion,(RegionsNum+1):(RegionsNum+NewRegionNum-1)]; % id of new-partitioned regions
        RegionSet(NewRegionId,:) = NewRegion;
        RegionSampleId(NewRegionId) = NewRegionSampleId;
        RegionsNum = RegionsNum+NewRegionNum-1;  % update total region number
        PartitionableNum = PartitionableNum-1+sum(NewRegion(:,2));  % update partitionable region number
        % update the maximal partition depth
        MaxPartitionDepth_temp = max(NewRegion(:,1));
        if MaxPartitionDepth_temp>CurrentMaxPartitionDepth
            CurrentMaxPartitionDepth = MaxPartitionDepth_temp;
            GroupInf(1:RegionsNum,4) = max(2,round(GroupInf(1:RegionsNum,1).*RegionSet(1:RegionsNum,2).*RegionSet(1:RegionsNum,1)./CurrentMaxPartitionDepth));
            sum_adj_size = sum(GroupInf(1:RegionsNum,4));
            BestGroupInf(3) = max(2,round(GroupInf(BestGroupInf(1),1)*RegionSet(BestGroupInf(1),1)/CurrentMaxPartitionDepth));
            all_weight_changed = 1;
        end
        % first stage sampling and group information updating
        for j = NewRegionId
            if RegionSet(j,2)==0
                %% non-partitionable region
                if isempty(RegionSampleId{j})
                    N0 = N0+1;
                    SampleSet(N0,2:end) = Sampling(1,RegionSet(j,3:end));
                    RegionSampleId{j} = N0;
                end
                GroupInf(j,1) = length(RegionSampleId{j});
                group_update(j)
                sum_adj_size = sum_adj_size-GroupInf(j,4);
                GroupInf(j,4) = 0;
                sum_weight = sum_weight - GroupInf(j,5);
                GroupInf(j,5) = 0;
                % all samples in this region are potential optima
                OptCandidate((OptCandidateNum+1):(OptCandidateNum+GroupInf(j,1)),1) = RegionSampleId{j};  % update the candidate optimal solution
                OptCandidateNum = OptCandidateNum+GroupInf(j,1);
                NewOptCandidate = 1;
            else
                %% partitionable region
                GroupInf(j,1) = length(RegionSampleId{j});  % update the sample size
                if GroupInf(j,1)>=MaxSampleSize
                    % if the sample size in the new region also reaches the threshold
                    PartitionListNum = PartitionListNum+1;
                    PartitionList(PartitionListNum) = j;
                else
                    if GroupInf(j,1)<n0
                    % if the sample size in the new region is not enough
                        Add_n = min(n0-GroupInf(j,1),N-N0);
                        NewSampleId = ((N0+1):(N0+Add_n))';
                        SampleSet(NewSampleId,2:end) = Sampling(Add_n,RegionSet(j,3:end));
                        N0 = N0+Add_n;
                        RegionSampleId{j} = [RegionSampleId{j};NewSampleId];
                        GroupInf(j,1) = GroupInf(j,1)+Add_n;
                    end
                    group_update(j)
                    WeightChangedListNum = WeightChangedListNum+1;
                    WeightChangedList(WeightChangedListNum) = j;
                end
            end
        end
        i = i+1;
    end
    PartitionListNum = 0;
%% update the optima set and do local search
    if NewOptCandidate == 1
        optima = extract();
        if ~isempty(OptSet)
            % remove the found optimum
            remove_opt(optima(:,3:end));
        end
        % statistc
        NumOpt_old = NumOpt;
        NumOpt = size(optima,1);
        if NumOpt_old ~= NumOpt
            NumIterationRepeat = 0;
        end
        NewOptCandidate = 0;
    end
    NumIterationRepeat = NumIterationRepeat+1;
%% update the weights using baqm
    if all_weight_changed == 1
    % if the maximal partition depth or the minimal quantile is changed
        WeightChangedListNum = PartitionableNum;
        WeightChangedList(1:PartitionableNum) = find(RegionSet(1:RegionsNum,2)==1)';
        all_weight_changed = 0;
    end
    baqm(WeightChangedList(1:WeightChangedListNum))
    WeightChangedListNum = 0;
%% additional budget
    Add_n = min([N-N0,Delta]);
    [ n2 ] = add_budget(Add_n);
    for i = 1:size(n2,1)
        iregion = n2(i,1);
        NewSampleId = ((N0+1):(N0+n2(i,2)))';
        SampleSet(NewSampleId,2:end) = Sampling(n2(i,2),RegionSet(iregion,3:end));
        N0 = N0+n2(i,2);
        RegionSampleId{iregion} = [RegionSampleId{iregion};NewSampleId];
        GroupInf(iregion,1) = GroupInf(iregion,1)+n2(i,2);
        if GroupInf(iregion,1)>=MaxSampleSize && RegionSet(iregion,2)==1
            PartitionListNum = PartitionListNum+1;
            PartitionList(PartitionListNum) = iregion;
        else
            group_update(iregion)
            WeightChangedListNum = WeightChangedListNum+1;
            WeightChangedList(WeightChangedListNum) = iregion;
        end
    end
end
%% final results
RegionSet((RegionsNum+1):end, :) = [];
SampleSet((N0+1):end, :) = [];
SampleSet(:,1)=[];
RegionSampleId((RegionsNum+1):end, :) = [];
optima = optima(:,2:end);

%% Plots for 2 dimension box constraint problems
% if d==2
%     figure()
%     hold on
%     scatter(SampleSet(:,3),SampleSet(:,4),2,SampleSet(:,2),'filled');
%     for i = 1:RegionsNum
%         rectangle('Position',[RegionSet(i,[3,5]),RegionSet(i,[4,6])-RegionSet(i,[3,5])])
%     end
%     for i = 1:size(optima,1)
%         plot(optima(:,3),optima(:,4),'x')
%     end
%     hold off
% end


%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% FUNCTIONS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function baqm(list)
    %baqm: calculate the weight for regions in the list using Budget
    %Allocation for Quantile Minimization
        if ~isempty(list)
            Wk_best = (1+norminv(alpha,0,1)^2-1/BestGroupInf(3));
            K = length(list); % number of groups
            cv_hat_inv = (GroupInf(list,2)-BestGroupInf(2))./GroupInf(list,3);
            Wk = (1+(cv_hat_inv.^2)-1./GroupInf(list,4));
            Ckb = Wk_best./Wk;
            P_kb = fcdf(Ckb, GroupInf(list,4)-1, ones(K,1).*BestGroupInf(3)-1); % the probability that cv_k>cv_b
            weights = P_kb./(1-P_kb); % Ni / Nb
            sum_weight = sum_weight-sum(GroupInf(list,5));
            GroupInf(list,5) = weights;
            sum_weight = sum_weight+sum(weights);
        end
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function [ n2 ] = add_budget(NewBudget)
    % add_budget: allocate NewBudget according to the weights
    % n2: K -by- 2 matrix, [region id, new sample size]
        N_tot = sum_adj_size+NewBudget;
        n = N_tot/sum_weight.*GroupInf(1:RegionsNum,5);
        n2 = zeros(RegionsNum,1);
        n = n-GroupInf(1:RegionsNum,4);
        while NewBudget>0
            [n_delta, k] = max(n);
            if n_delta<1
                add = 1;
            else
                add = min(NewBudget,round(n_delta));
            end
            n2(k) = n2(k)+add;
            n(k) = n(k)-add;
            NewBudget = NewBudget-add;
        end
        n2_id = find(n2>0);
        n2 = [n2_id, n2(n2_id,1)];
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function group_update( groupId )
    % update the group information
        Samplevalue_temp = SampleSet(RegionSampleId{groupId},2);
        % update group information
        GroupInf(groupId,2) = mean(Samplevalue_temp);  % group mean
        GroupInf(groupId,3) = sqrt(var(Samplevalue_temp));  % group variance
        sum_adj_size = sum_adj_size-GroupInf(groupId,4);
        GroupInf(groupId,4) = max([2,round(GroupInf(groupId,1)*RegionSet(groupId,1)/CurrentMaxPartitionDepth)]);  % modified sample size
        sum_adj_size = sum_adj_size+GroupInf(groupId,4);
        new_quantile = GroupInf(groupId,2)+GroupInf(groupId,3)*norminv(alpha,0,1);
        if groupId==BestGroupInf(1) || new_quantile<BestGroupInf(2)
            if new_quantile<=BestGroupInf(2)
                BestGroupInf = [groupId, new_quantile, GroupInf(groupId,4)];
            else
                [BestGroupInf(2),BestGroupInf(1)] = min(GroupInf(1:RegionsNum,2)+GroupInf(1:RegionsNum,3).*norminv(alpha,0,1));
                BestGroupInf(3) = max([2,round(GroupInf(BestGroupInf(1),1)*RegionSet(BestGroupInf(1),1)/CurrentMaxPartitionDepth)]);
            end
            all_weight_changed = 1;
        end
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function [ Optima ] = extract()
    %Nitching: obtain the final optima set
    % Optima: [candidate id, sample id, obj, x1, x2, ...]
        %% sort new potential optima
        NewCandidateId_f = ((OptCandidate_checked+1):OptCandidateNum)';
        [~,sort_id] = sort(SampleSet(OptCandidate(NewCandidateId_f,1),2));
        OptCandidate(NewCandidateId_f,:) = OptCandidate(sort_id+OptCandidate_checked,:);
        NewSampleId_f = OptCandidate(NewCandidateId_f,1);
        NewCandidateSize_f = OptCandidateNum-OptCandidate_checked;
        %% clearing
        candidate_finished = 0;
        while 1
            if candidate_finished<size(optima,1)
                % for previous optima
                candidate_id = optima(candidate_finished+1,1);
                sampleid = optima(candidate_finished+1,2);
            else
                % for new optima candidate
                if candidate_finished==size(optima,1)
                    candidate_id = OptCandidate_checked;
                    next = find(OptCandidate(NewCandidateId_f,2),1);
                else
                    next = find(OptCandidate((candidate_id+1):OptCandidateNum,2),1);
                end
                if isempty(next) || (candidate_id+next)>OptCandidateNum
                    break
                else
                    candidate_id = candidate_id + next;
                    sampleid = OptCandidate(candidate_id,1);
                end
            end
            if candidate_id<=OptCandidate_checked
                Distance = (sqrt(sum((repmat(SampleSet(sampleid,3:end),[NewCandidateSize_f,1])-SampleSet(NewSampleId_f,3:end)).^2,2)))';
                NeighborCandidateId = NewCandidateId_f(Distance<=radius);
            else
                Distance = (sqrt(sum((repmat(SampleSet(sampleid,3:end),[OptCandidateNum,1])-SampleSet(OptCandidate(1:OptCandidateNum,1),3:end)).^2,2)))';
                Distance(candidate_id) = inf;
                NeighborCandidateId = find(Distance<=radius);
            end
            NeighborSampleId = OptCandidate(NeighborCandidateId,1);
            WorseCandidateId = NeighborCandidateId(SampleSet(NeighborSampleId,2)>SampleSet(sampleid,2));
            OptCandidate(WorseCandidateId,2) = 0; % delete all worse neighbor
            if SampleSet(sampleid,2)>min(SampleSet(NeighborSampleId,2))
                OptCandidate(candidate_id,2) = 0;
            end
            candidate_finished = candidate_finished+1;
        end
        optima_temp = find(OptCandidate(1:OptCandidateNum,2)==1);
        Optima = [optima_temp, SampleSet(OptCandidate(optima_temp,1),:)];
        OptCandidate_checked = OptCandidateNum;
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function remove_opt( new_samples )
    % remove_opt: remove the found optimal solution from OptSetRemain
        new_samples(new_samples(:,1)>=(OptSet.fval+OptSet.epsilon),:) = [];
        for ii = 1:size(new_samples,1)
            OptDis = sqrt(sum((repmat(new_samples(ii,2:end),[size(OptSetRemain,1),1]) - OptSetRemain).^2,2));
            [MinDis, OptId] = min(OptDis);
            if MinDis<OptSet.radius
                OptSetRemain(OptId,:) = [];
            end
        end
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
end

